import torch
import torch.nn as nn

class kAE(nn.Module):

	def __init__(self, in_sz, hd_sz, out_cls):
		
		super(kAE, self).__init__()
		self.out_cls = out_cls
		self.in_sz = in_sz
		base = hd_sz

		self.AE = nn.ModuleList([nn.Sequential(
			nn.Linear(in_sz, base * 2),
			nn.ReLU(), 
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, base),
			nn.ReLU(),
			nn.BatchNorm1d(base),
			nn.Linear(base, base * 2),
			nn.ReLU(),
			nn.BatchNorm1d(base * 2),
			nn.Linear(base * 2, in_sz),
			nn.ReLU()
			) for i in range(out_cls)])

	def forward(self, x):

		return torch.cat([self.AE[i](x).unsqueeze(1) for i in range(self.out_cls)], 1)
